import matplotlib.pyplot as plt

decisionNode = dict(boxstyle="square,pad=0.5", fc="0.9")
leafNode = dict(boxstyle="round4, pad=0.5", fc="0.9")
arrow_args = dict(arrowstyle="<-", connectionstyle="arc3", shrinkA=0,
                  shrinkB=16)


def plot_node(node, center, parent, type):
    draw_tree.ax1.annotate(node, xy=parent, xycoords='axes fraction',
                            xytext=center, textcoords='axes fraction', va="top",
                            ha="center", bbox=type, arrowprops=arrow_args)


def count_leaves(tree):
    num = 0
    first = list(tree.keys())[0]
    second = tree[first]
    for key in second.keys():
        if type(second[key]).__name__ == 'dict':
            num += count_leaves(second[key])
        else:
            num += 1
    return num


def get_depth(tree):
    depth = 0
    first = list(tree.keys())[0]
    second = tree[first]
    for key in second.keys():
        if type(second[key]).__name__ == 'dict':
            current = 1 + get_depth(second[key])
        else:
            current = 1
        if current > depth:
            depth = current
    return depth


def plot_txt(cnt, parent, txt):
    x = (parent[0] - cnt[0]) / 2.0 + cnt[0]
    y = (parent[1] - cnt[1]) / 2.0 + cnt[1]
    draw_tree.ax1.text(x, y, txt)


def plot_tree(tree, parent, node):
    num = count_leaves(tree)
    depth = get_depth(tree)
    first = list(tree.keys())[0]
    cnt = (plot_tree.x_off + (1.0 + float(num)) / 2.0 / plot_tree.W,
           plot_tree.y_off)
    plot_txt(cnt, parent, node)
    plot_node(first, cnt, parent, decisionNode)
    second = tree[first]
    plot_tree.y_off = plot_tree.y_off - 1.0 / plot_tree.D
    for key in second.keys():
        if type(second[key]).__name__ == 'dict':
            plot_tree(second[key], cnt, str(key))
        else:
            plot_tree.x_off = plot_tree.x_off + 1.0 / plot_tree.W
            plot_node(second[key], (plot_tree.x_off, plot_tree.y_off),
                      cnt, leafNode)
            plot_txt((plot_tree.x_off, plot_tree.y_off), cnt, str(key))
    plot_tree.y_off = plot_tree.y_off + 1.0 / plot_tree.D


def draw_tree(tree):
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 显示中文标签
    fig = plt.figure(1, facecolor='white')
    fig.clf()
    axprops = dict(xticks=[], yticks=[])
    draw_tree.ax1 = plt.subplot(111, frameon=False, **axprops)
    plot_tree.W = float(count_leaves(tree))
    plot_tree.D = float(get_depth(tree))
    plot_tree.x_off = -0.5 / plot_tree.W
    plot_tree.y_off = 1.0
    plot_tree(tree, (0.5, 1.0), '')
    plt.show()
